package com.demo.java.limiter;

import com.demo.java.utils.LogUtils;
import com.demo.java.utils.ThreadUtils;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 令牌桶限流
 *
 * 桶的容量固定，发放令牌速度固定，主要判断桶内剩余令牌数量是否大于请求数量，大于则放行，否则请求失败
 *
 * 可以通过提升令牌发放速度来提升后端处理能力。
 */
public class TokenBucketLimiterTest {

    public static void main(String[] args) throws InterruptedException {
        TokenBucketLimiter limiter = new TokenBucketLimiter();
        ExecutorService executor = Executors.newFixedThreadPool(10);
        int threads = 2;
        int loop = 20;
        AtomicInteger limited = new AtomicInteger(0);
        long start = System.currentTimeMillis();

        CountDownLatch latch = new CountDownLatch(threads);


        for (int i = 0; i < threads; i++) {
            executor.submit(() -> {
                for (int j = 0; j < loop; j++) {
                    if (!limiter.tryAcquire(1)) {
                        limited.incrementAndGet();
                    }
                    ThreadUtils.sleep(200, TimeUnit.MILLISECONDS);
                }
                System.out.println(LogUtils.logt("执行结束"));
                latch.countDown();
            });
        }


        System.out.println(LogUtils.log("等待结束"));
        latch.await();

        long cost = (System.currentTimeMillis() - start)/1000;
        System.out.println(LogUtils.log("被限制次数: " + limited.get()) + ", 通过的次数: " + (threads * loop - limited.get()));
        System.out.println(LogUtils.log("限制比例: " + (limited.get()/(float)(threads * loop))));
        System.out.println(LogUtils.log("花费时间: " + cost));

        executor.shutdown();
    }

    static class TokenBucketLimiter {
        // 最后一次发放令牌的时间
        private long lastTime = System.currentTimeMillis();
        // 令牌发放速度
        private int rate = 2;
        // 令牌桶的容量
        private int capacity = 2;
        // 桶内剩余令牌数量
        private AtomicInteger tokens = new AtomicInteger();

        public boolean tryAcquire(int tryCount) {
            // 当前时间
            long now = System.currentTimeMillis();
            // 计算时间单位内发放的令牌数量
            int permitsCount = (int)(((now - lastTime)/1000)*rate);
            // 总的令牌数= 剩余令牌数 + 单位时间发放的令牌数
            int allPermits = tokens.get() + permitsCount;
            // 令牌数整形，令牌个数最大不能超过桶的容量
            tokens.set(Math.min(capacity, allPermits));
            System.out.println(LogUtils.log("capacity: " + capacity + ", tokens: " + tokens.get() + ", permitsCount: " + permitsCount));

            // 判断剩余令牌数是否大于请求数，大于才放过请求，否则拒绝
            if (tokens.get() > tryCount) {
                tokens.addAndGet(- tryCount);
                lastTime = now;
                return true;
            }

            return false;
        }

    }
}
